import math

from matrix import *


class ConvolvedFeatureLayer:
    """
    卷积特征层类。
    """

    maps = None

    def __init__(self, num):
        self.maps = [1] * num

    def handler(self, data, kernel_width, kernel_height):
        # 计算公式：(图片矩阵列数 - 卷积核列数) / 步长 + 1
        feature_map_width = data.shape[2] - kernel_width + 1
        # 计算公式：(图片矩阵行数 - 卷积核行数) / 步长 + 1
        feature_map_height = data.shape[1] - kernel_height + 1
        # 深度
        depth = data.shape[0]

        for index in range(len(self.maps)):
            feature_map = ConvolvedFeatureMap(feature_map_width, feature_map_height, kernel_width, kernel_height, depth)
            feature_map.build_map(data)
            self.maps[index] = get_matrix(feature_map.map)


class ConvolvedFeatureMap:
    """
    卷积特征图类。
    """

    # 卷积特征图，其实就是一个矩阵。
    map = None

    # 卷积核，其实就是一个权值矩阵。
    kernel = None

    # 步数
    step = 1

    def __init__(self, width, height, kernel_width, kernel_height, kernel_depth=1):
        """
        初始化卷积图。
        :param width: 卷积图的宽度。
        :param height: 卷积图的高度。
        :param kernel_width: 卷积核的宽度。
        :param kernel_height: 卷积核的高度。
        :param kernel_depth: 卷积核的深度。
        :return:
        """
        # 初始化卷积特征矩阵，每张矩阵只有一层。
        self.map = np.zeros((height, width))

        # 初始化 kernel_depth 层的卷积核（宽：kernel_width，高：kernel_height）。
        self.kernel = np.random.randn(kernel_depth, kernel_height, kernel_width) / math.sqrt(
            2.0 / kernel_width * kernel_height)

    def set_kernel(self, kernel):
        self.kernel = kernel

    def get_convolved_feature(self, data, x, y):
        """
        提取卷积特征。
        卷积特征 = Relu(SUM(小图矩阵 X 卷积核（即：权重矩阵）))
        这里的计算流程和全连接神经网络是一样的，但有一点不同的就是：
            全连接神经网络是矩阵的积在求和。
            卷积特征是矩阵中每个对应元素的积在求和。
        如果传入的是彩色图的矩阵的话，由于有 R, G, B 三个通道，所有就有对应的三个矩阵。
        这时卷积核也必须要有对应的 3 个。
        原理只要看这张图就能明白：http://ufldl.stanford.edu/wiki/index.php/%E5%8D%B7%E7%A7%AF%E7%89%B9%E5%BE%81%E6%8F%90%E5%8F%96
        :param data: 特征矩阵，该样本矩阵的深度必须与初始化的卷积核的深度一致。
        :param x: 横坐标位置。
        :param y: 纵坐标位置。
        :return: 卷积特征。
        """

        # 在特征矩阵中指定的位置裁剪出与卷积核相同矩阵（即：将一张大图拆分在多张小图分别处理来减少计算量）。
        matrix = slice_matrix(data, x, y, self.kernel.shape[2], self.kernel.shape[1])

        # 初始化输出值。
        output = 0

        # 遍历每一层的特征矩阵与卷积核，将每一层的计算结果合并。
        for i in range(self.kernel.shape[0]):
            # 获得当前层的特征矩阵。
            a1 = get_matrix(matrix, i)
            # 获得当前层的卷积核，注意：每一个完整的特征矩阵对应的卷积核是相同的。
            w1 = get_matrix(self.kernel, i)
            # 将当前层的特征矩阵与卷积核进行点积运算并求和。
            v1 = (a1 * w1).sum()
            # 将每一层的计算结果求和。
            output += v1

        # 对上一步的计算结果使用 Relu 函数进行运算。
        if output < 0:
            output = 0

        # 卷积特征。
        return output

    def build_map(self, data):
        """
        生成卷积特征图。
        :param data: 特征矩阵。
        :return:
        """
        if data.shape[0] != self.kernel.shape[0]:
            raise ValueError("传入的特征矩阵的层数与卷积核的层数不一致")

        # 将数组转换成矩阵。
        data = np.array(data)
        max_x = data.shape[1]
        max_y = data.shape[2]

        # 定义卷积特征图的行数（高度）。
        row = self.map.shape[1]
        # 定义卷积特征图的列数（宽度）。
        col = self.map.shape[0]
        x = 0
        y = 0

        # 将完整的矩阵 data 中选取其中一小块，学习其中的特征，使用相同的方法遍历整个矩阵。
        for r in range(row):
            for c in range(col):
                # 学习从坐标（x, y）开始的矩阵的特征。
                convolved_feature = self.get_convolved_feature(data, x, y)
                # 保存学习到的特征到卷积特征矩阵里。
                self.map[r, c] = convolved_feature
                if x < max_x:
                    # 位置往右移动 step 步。
                    x += self.step
            x = 0
            if y < max_y:
                # 位置往下移动 step 步。
                y += self.step

        # 卷积特征矩阵。
        return self.map


def max_pooling_handler(maps):
    """
    池化。
    :param maps: 卷积特征图列表。
    :return: 池化后的矩阵。
    """
    data = np.array(maps)
    row = int(data.shape[2] / 2)
    col = int(data.shape[1] / 2)

    # 创建一个 row 行，col 列的二维数组。
    convolved_max_pool = np.zeros((data.shape[0], row, col))

    for p in range(data.shape[0]):
        x = 0
        y = 0
        for r in range(row):
            for c in range(col):
                matrix = slice_matrix(data, x, y, 2, 2, p)
                max_index = np.argmax(matrix)

                if max_index > 1:
                    row_max_index = 1
                    col_max_index = max_index - 2
                else:
                    row_max_index = 0
                    col_max_index = max_index

                convolved_max_pool[p, r, c] = matrix[row_max_index, col_max_index]
                if x < data.shape[1]:
                    x += 2
            x = 0
            if y < data.shape[2]:
                y += 2

    return convolved_max_pool
